from camel.models import ModelFactory
from camel.types import ModelPlatformType, ModelType
from camel.loaders import UnstructuredIO
from camel.storages import Neo4jGraph
from camel.retrievers import AutoRetriever
from camel.types import StorageType
from camel.agents import ChatAgent, KnowledgeGraphAgent
from camel.messages import BaseMessage
from camel.embeddings import SentenceTransformerEncoder
from py2neo import Graph


# 设置检索器
camel_retriever = AutoRetriever(
    vector_storage_local_path="local_data/embedding_storage",
    storage_type=StorageType.QDRANT,
    embedding_model=SentenceTransformerEncoder(model_name='intfloat/e5-large-v2'),
)

from dotenv import load_dotenv
import os

load_dotenv()
Model_Type = os.getenv("MODEL_TYPE")
Model_Api = os.getenv("ZHIPU_API_KEY")
Model_Url = os.getenv("MODEL_URL")
model = ModelFactory.create(
    model_platform=ModelPlatformType.OPENAI,
    model_type=Model_Type,
    api_key=Model_Api,
    url=Model_Url
)

n4j = Neo4jGraph(
    url='bolt://localhost:7687',
    username='neo4j',
    password='Shao123.',
)


def delete_neo4j_graph():
    graph = Graph(
        uri='bolt://localhost:7687',
        auth=('neo4j', 'Shao123.'),
    )
    graph.delete_all()


uio = UnstructuredIO()
kg_agent = KnowledgeGraphAgent(model=model)

with open('ceshi.txt', 'r', encoding='utf-8') as f:
    txt_content = f.read()
query = "Python 程序在运行时出现错误，错误信息为：'TypeError: unsupported operand type (s) for +: 'int' and'str''。"
vector_result = camel_retriever.run_vector_retriever(
    query=query,
    contents=txt_content,
)

eles = uio.create_element_from_text(
    text=txt_content,
)

graph_element = kg_agent.run(eles, parse_graph_elements=True)

n4j.add_graph_elements(graph_elements=[graph_element])

query_element = uio.create_element_from_text(
    text=query,
    element_id='1'
)

# 让知识图谱agent从查询中提取节点和信息

ans_element = kg_agent.run(query_element, parse_graph_elements=True)

# 匹配从query中获得的实体在知识图谱存储内容中的信息
kg_result = []
for node in ans_element.nodes:
    n4j_query = f"""
MATCH (n {{id: '{node.id}'}})-[r]->(m)
RETURN 'Node ' + n.id + ' (label: ' + labels(n)[0] + ') has relationship ' + type(r) + ' with Node ' + m.id + ' (label: ' + labels(m)[0] + ')' AS Description
UNION
MATCH (n)<-[r]-(m {{id: '{node.id}'}})
RETURN 'Node ' + m.id + ' (label: ' + labels(m)[0] + ') has relationship ' + type(r) + ' with Node ' + n.id + ' (label: ' + labels(n)[0] + ')' AS Description
"""
    result = n4j.query(query=n4j_query)
    kg_result.extend(result)

kg_result = [item['Description'] for item in kg_result]

# 合并来自向量搜索和知识图谱实体搜索的结果
comined_results = str(vector_result) + "\n".join(kg_result)

sys_msg = BaseMessage.make_assistant_message(
    role_name='CAMEL_Agent',
    content="""
    你是一个有用的助手来回答问题，
    我会给你提供原始的查询和检索到的上下文，
    根据检索到的上下文回答原始查询
    """
)

camel_agent = ChatAgent(
    system_message=sys_msg,
    model=model
)

user_prompt = f"""
原始查询是{query},
检索到的上下文是{comined_results}
"""

user_msg = BaseMessage.make_user_message(
    role_name="User",
    content=user_prompt
)
response = camel_agent.step(user_msg)
print(response.msgs[0].content)
